#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   TcpNameGet.py
@Time    :   2024-11-27 15:30
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   TCP端口服务名获取
"""
from argparse import ArgumentParser
import os
import nmap
from loguru import logger


class TcpNameGet:
	def __init__(self, ip, mask=24, port="1-65535", service='ssh'):
		"""
		初始化网络扫描器

		:param ip: 设备的IP地址
		:param mask: 子网掩码，默认为24
		:param port: 要扫描的端口范围，可以使用格式: '22,53,110,143-4564', 默认 1-65535
		:param service: 要扫描的服务名称，默认为 "ssh"
		"""
		self.service = service
		self.ip = ip
		self.mask = mask

		# 创建 self.portscanner 对象
		self.nm = nmap.PortScanner()

		# 定义要扫描的网络段
		self.network = f"{ip}/{mask}"

		# 定义要扫描的端口，SSH 默认端口是 22
		self.ports = str(port)

	def check_cmd_path(self):
		"""
		检查本机是否安装nmap命令
		:return:
		"""
		if os.system("nmap -h >/dev/null 2>&1") != 0:
			logger.error("nmap 未安装，请安装 nmap")
			exit(1)
		logger.info("nmap 已安装")

	def scan(self):
		"""
		将开放的SSH扫描结果写入文件。

		此方法遍历扫描结果中的所有主机，检查指定端口是否开放。
		如果开放，则写入主机信息和SSH版本信息到文件；如果没有开放，则写入相应信息。
		"""
		# 扫描网络
		self.nm.scan(hosts=str(self.network), ports=str(self.ports), arguments='-sV')  # -sV 用于服务版本检测
		open_file = open(f"TcpNameGet_open_{self.ip}-{self.mask}.txt", "w")
		if len(self.nm.all_hosts()) == 0:
			logger.error("没有扫描到任何主机")
			exit(2)
		# 保存扫描结果到文件
		for host in self.nm.all_hosts():
			# 获取主机IP地址
			# host_ip = self.nm[host].hostname()
			logger.info(f"Host: ({host})")
			tcp_info = self.nm[host]['tcp']
			if len(tcp_info) == 0:
				logger.warning(f"{host}: 没有开放端口")
				continue
			logger.info("检测信息如下")
			for port_, info in tcp_info.items():
				service = ''
				print(f"{port_}: {info}")
				try:
					service = str(tcp_info[port_]['product']).lower()
				except Exception as e:
					logger.error(e)
					exit(4)

				if service == self.service.lower():
					logger.info(f"{self.service} - {host}: {port_} is open")
					try:
						open_file.write(f"{host}-{port_}-{self.service.lower()}\n")
					except Exception as e:
						logger.error(e)
						exit(3)

	def start(self):
		self.check_cmd_path()
		self.scan()


if __name__ == '__main__':
	# 传参参数及说明
	arg = ArgumentParser(description='Tcp服务端口扫描,通过设置指定服务名称获取该服务开启的所有端口', prog="TcpNameGet")
	arg.add_argument("-i", "--ip", help="IP地址", type=str, required=True)
	arg.add_argument("-m", "--mask", help="子网掩码", type=int, default=24)
	h = "指定要扫描的端口，可以使用格式: '22,53,110,143-4564', 默认 1-65535"
	arg.add_argument("-p", "--port", help=h, type=str, default='1-65535')
	arg.add_argument("-n", "--name", help="设置服务名称,默认: OpenSSH", type=str, default="OpenSSH")
	args = arg.parse_args()
	logger.info(args)
	logger.info("开始扫描")
	ssh = TcpNameGet(args.ip, args.mask, args.port, service=args.name)
	ssh.start()
